from data_utils import *
from bayes_classifier import *
from utils import *
from sklearn.naive_bayes import GaussianNB

if __name__ == '__main__':
    # 加载数据集
    watermelon_dataset = WatermelonDataset('./watermelon_dataset.txt')
    # 划分训练集和测试集
    # 此处由于我们需要预测后面的新样本，所以将所有数据作为训练数据
    train_data, test_data, train_target, test_target = watermelon_dataset.get_train_test_set(1.0)
    # 实例化分类器
    my_classifier = NaiveBayesClassifier()
    # 拟合数据
    my_classifier.fit(train_data, train_target)

    # 准备一个新样本
    new_data = [['青绿', '蜷缩', '浊响', '清晰', '凹陷', '硬滑', 0.697, 0.460]]
    # 转换为数据集格式
    input_data = watermelon_dataset.convert_new_data(new_data)
    # 预测类别
    predict = my_classifier.predict(input_data)[0]
    print('my model predict:', predict)
    print('是否好瓜：', reverse_dict(watermelon_dataset.str_map['好瓜'])[predict])

    # 下面使用sklearn中的实现进行预测，以进行比较
    gnb = GaussianNB()
    # 训练模型
    gnb.fit(train_data, list(train_target.reshape(-1)))
    # 预测
    sklearn_pred = gnb.predict(input_data)[0]
    print('sklearn model predict:', sklearn_pred)
    print('是否好瓜：', reverse_dict(watermelon_dataset.str_map['好瓜'])[sklearn_pred])


